Skip to content

Add fp8_e3m4 dtype support - #97

Merged
SmoothThunk merged 9 commits into
leanprover:mainfrom
SmoothThunk:float8-e3m4
Jul 28, 2026
Merged

Add fp8_e3m4 dtype support#97
SmoothThunk merged 9 commits into
leanprover:mainfrom
SmoothThunk:float8-e3m4

Conversation

@SmoothThunk

@SmoothThunk SmoothThunk commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

fp8_e3m4:

  • Contains +-inf and NaN
  • Float.lean: decode and encode functions use round-to-nearest-even
  • Dtype.lean: all match arms (arithmetic, casts, promotion, lossless)
  • Npy.lean: maps to V1 (same as e4m3, ambiguous in npy format)
  • Test.lean: IO tests for decode, arithmetic, cast, and overflow

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Code Review

Single-pass review, ranked most-severe first.

1. TensorLib/Tensor.lean:604 — e3m4 missing from toFloat32Tree / toFloat64Tree

toFloat32Tree and toFloat64Tree have no arm for .float8_e3m4, so 1-byte e3m4 elements fall through to Float32.ofLEByteArray, which requires 4 bytes and errors on every element.

Failure scenario: arr.dtype = .float8_e3m4; toFloat32Tree arr (or the ! variant) hits the _ arm and calls Float32.ofLEByteArray on a 1-byte ByteArrayErr (or panic via get!). e3m4 tensors are effectively undecodable through the standard tree accessors this PR otherwise extends.

2. TensorLib/Dtype.lean:190join not commutative between e3m4 and e4m3

joinOrdered adds (.float8_e3m4, .float8_e4m3) → float8_e3m4, but the earlier .float8_e4m3 arm falls through to .float8_e4m3, _ => none, and both types have equal itemsize so join() doesn't swap — the promotion is not commutative.

Failure scenario: Dtype.join .float8_e3m4 .float8_e4m3 = some float8_e3m4, but Dtype.join .float8_e4m3 .float8_e3m4 = none. The join-commutativity PBT (line 1551) misses this only because gen (line 55) never samples float8_e3m4.

3. TensorLib/Dtype.lean:307lossless(e3m4, e4m3) is incorrectly true

lossless marks .float8_e3m4 → .float8_e4m3 as true, but e4m3 has only 3 mantissa bits vs e3m4's 4 — a mantissa bit is always dropped, so this cast is lossy.

Failure scenario: e3m4 value 1.0625 (1.0001₂ · 2⁰) has no exact e4m3 representation and rounds to 1.0 or 1.125. lossless returning true here also violates losslessAntiSymmetric semantics for the pair and misleads any caller that gates precision-preserving conversions on this predicate.

4. TensorLib/Dtype.lean:55gen : Gen Dtype not updated

gen was not updated to include float8_e3m4, so every PBT that samples Dtype (join commutativity in this same file, plus any downstream generators) has zero coverage of the new type.

Failure scenario: The example (a b : Dtype) : Dtype.join a b == Dtype.join b a PBT (line 1551) reports Unable to find a counter-example even though finding #2 proves it is false, because a/b are never drawn as float8_e3m4.

5. TensorLib/Dtype.lean:310lossless missing e3m4 → fp16 / bf16

lossless is missing .float8_e3m4 → .float16 and .float8_e3m4 → .bfloat16, both of which are actually lossless (fp16/bf16 comfortably contain e3m4's range and mantissa).

Failure scenario: Any caller using lossless to decide whether to promote e3m4 into fp16/bf16 for accumulation will wrongly refuse the cast, inconsistent with the e4m3/e5m2 rows just above which list both targets.

6. TensorLib/Npy.lean:117 — e3m4 round-trips through .npy as e4m3

dtypeNameToNpyString emits "V1" for both e4m3 and e3m4 while fromNpyString maps "<V1" only to e4m3, so writing an e3m4 tensor and reading it back silently reinterprets the bytes as e4m3.

Failure scenario: Tensor.toNpy on an e3m4 tensor writes header dtype "<V1"; reading the file yields a tensor with dtype float8_e4m3 whose bytes now decode to entirely different Float32 values. The comment acknowledges the ambiguity but no roundtrip guard or explicit error prevents silent data corruption.

7. TensorLib/Dtype.lean:187 — style nit: bare bool in or-pattern

This or-pattern writes bool without the leading dot while every neighbor in joinOrdered uses .bool (lines 168, 181, 197, 203). Inside namespace Dtype it resolves to the constructor, but the inconsistency invites confusion with a pattern-variable binding.

Failure scenario: Readability regression — a maintainer copy-pasting this arm into a scope where Dtype is not opened would silently turn bool into a wildcard variable binding matching every dtype.


Caveat: single-pass review, no adversarial verify or multi-agent fan-out ran. Findings 3 and 4 in particular are worth spot-checking.

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Fixes:

  • Fixed-- added fp8_e3m4 arms to both toFloat32Tree and toFloat64Tree.
  • Added fp8_e3m4 arms to toFloat32Tree/toFloat64Tree and added | .float8_e4m3, .float8_e3m4 => float8_e3m4 before the wildcard in joinOrdered. Also added fp8_e3m4 case to the gen.
  • Removed float8_e4m3 from the lossless cases — e3m4 has 4 mantissa bits vs e4m3's 3, so the cast is lossy. Added a comment noting the divergence from np.can_cast which returns true.
  • Added fp8_e3m4 to gen as noted before.
  • Np.can_cast returns False for all fp8 -> fp16/bf16 (e4m3, e5m2, and e3m4). The existing e4m3/e5m2 lossless entries for fp16/bf16 were already diverging from numpy so I fixed them to match numpy. Only fp32 and fp64 are lossless targets for fp8.
  • Added #guards documenting that both e3m4 and e4m3 serialize to "V1". The npy format cannot distinguish them — loads <V1 as raw bytes, not as either fp8 type. The limitation is now explicitly tested. A future improvement could error on e3m4 write or add an explicit dtype parameter to the reader.
  • Fixed typo

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Code review findings

Ranked most severe first.


1. join(float8_e4m3, float8_e3m4) = float8_e3m4 truncates e4m3's rangeTensorLib/Dtype.lean:192

Given tensor A of dtype e4m3 with value 100.0 (representable — e4m3 max is 448) and tensor B of dtype e3m4 with value 0.5, a ufunc calling x.dtype.join y.dtype (see Ufunc.lean:302) picks e3m4. Storing 100.0 in e3m4 (max 15.5) overflows to +inf, so add(A, B) silently produces +inf. By analogy, join(float16, bfloat16) in this same file returns none because neither dominates; e4m3 vs e3m4 has the same asymmetry (e4m3 has larger range, e3m4 has more mantissa), so none (or promotion to float32) is the correct answer.

2. lossless regression: e4m3/e5m2 → float16/bfloat16 flipped from true to falseTensorLib/Dtype.lean:297-304

The diff comments out .float8_e4m3, .float16, .float8_e4m3, .bfloat16, .float8_e5m2, .float16, and .float8_e5m2, .bfloat16 from the lossless table. But float16/bfloat16 both strictly dominate fp8_e4m3 and fp8_e5m2 (more mantissa bits, wider exponent range) — every finite e4m3/e5m2 value round-trips exactly. A caller that guards a widening cast with if lossless fromDtype toDtype then ... (e.g., the PBT at Dtype.lean:1474) will now take a lossy path for these exact widenings.

3. lossless missing arms for e3m4 → float16/bfloat16TensorLib/Dtype.lean:308-312

Same shape as finding 2: fp16 (mantissa=10) and bf16 both strictly dominate e3m4 (mantissa=4, exp range narrower), so the widening is exact, but the table reports it as lossy. Every finite e3m4 value round-trips through float16 exactly.

4. Pattern arm | .float8_e3m4, bool uses bool without a leading dotTensorLib/Dtype.lean:189

Every other Dtype constructor in this match is written with a dot (.bool, .int8, .uint8, .float8_e4m3). This compiles today only because namespace Dtype is open and Lean resolves the bare identifier bool to the constructor. If a future import or refactor introduces another bool in scope, this pattern silently degrades into a variable-binding pattern that matches everything, making join .float8_e3m4 anyType = float8_e3m4 and skipping all later arms. Latent hazard — fix by writing .bool.

5. Duplicate/misleading #guard in the npy testsTensorLib/Npy.lean:442

Line 442 is a verbatim copy of line 437 (fromNpyString \"<V1\" == .ok { name := .float8_e4m3, ... }). The preceding comment on line 441 says "e3m4 cannot round-trip through npy (reads back as e4m3)" — but the guard doesn't actually demonstrate that (e.g., a fromNpyString \"<V1\" != .ok e3m4 check would). One wasted assertion; reader confusion.

6. Misleading comment on the e3m4/e4m3 join armTensorLib/Dtype.lean:192

The comment "e3m4 absorbs e4m3 (since e3m4 has 4 mantissa bits vs e4m3's 3)" justifies the (incorrect) absorption by mantissa width alone, ignoring that e4m3's exponent range (up to 448) is far wider than e3m4's (up to 15.5). A future maintainer who applies this reasoning to another pair may reproduce the range-truncation bug from finding 1.

7. maxSafeNatForFloat8e3m4 := 15 breaks the shared formula without a commentTensorLib/Float.lean:45

Every other fp format uses Nat.pow 2 (mantissaBits + 1). For e3m4 the formula gives 32, but the actual max representable integer is 15 (since e3m4 max is 15.5). The hardcoded 15 is correct, but the deviation has no comment. A future change that adds another narrow-range fp format may blindly copy the formula and miss the range constraint, or someone may "normalize" this call site to the formula and break canCastFromNat for e3m4.


Findings 1-3 are the correctness bugs; 4 is a latent hazard; 5-7 are cleanup.

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Fixes:

  • Fixed; Changed join(e4m3, e3m4) to none in both directions. e4m3 has more range (max 448) while e3m4 has more mantissa (4 bits). numpy's result_type says e3m4 here but that causes silent overflow (100.0 → inf), so we diverge.
  • Added fp16/bf16 back as lossless targets for all three fp8 types — the casts are mathematically lossless (all values round-trip exactly). numpy's can_cast returns false here because ml_dtypes registers fp8 types as void types not because the cast is actually lossy. We diverge from numpy here.
  • Fixed in previous comment.
  • Added .bool
  • Replaced original guard with a guard that demonstrates the limitation.
  • Already addressed — the absorption was replaced with none.

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

I read the full diff, then the enclosing functions and adjacent fp8 (e4m3/e5m2) paths in Dtype.lean, Float.lean, Npy.lean, Tensor.lean, and Test.lean, and traced encoder/decoder rounding cases and cast graph coverage. No genuine correctness bugs found; findings below are mostly documented footguns and a comment placement issue.


Finding 1 — TensorLib/Npy.lean:117 — e3m4/e4m3 npy round-trip silently corrupts data

e3m4 and e4m3 both serialize to "V1" in npy with no distinguishing marker, so a save/load round-trip silently reinterprets e3m4 bytes as e4m3.

Failure scenario: user creates a Tensor with dtype .float8_e3m4 containing 1.0 (encoded bits 0x30) and calls Tensor.save! then Tensor.load. dtypeNameToNpyString writes "<V1" to the header; fromNpyString on read maps "<V1" to .float8_e4m3 (line 135). The bytes are unchanged, but e4m3 interprets 0x30 as 0.5 (exp=6, mant=0 -> 2^(6-7)). Numeric data is silently corrupted with no error surfaced to the caller, and downstream arithmetic uses the wrong dtype path. The diff acknowledges this as a "known limitation" but it qualifies as a real bug because save+load is expected to round-trip. Consider erroring on save for e3m4, or using a distinct marker.


Finding 2 — TensorLib/Dtype.lean:313 — misplaced lossless comment

The comment about e3m4->e4m3 lossiness is placed inside the e3m4-lossless list, between the bfloat16 and float32 arms, making it read as a note about float32.

The block reads:

| .float8_e3m4, .bfloat16
-- Note: np.can_cast(e3m4, e4m3) returns true but this is incorrect since precision is lost
| .float8_e3m4, .float32

A future maintainer will parse the comment as annotating the .float32 arm (or the preceding .bfloat16 arm), not as an explanation for why .float8_e4m3 is deliberately absent from the lossless list. Move the note above the block or next to the fallthrough | .float8_e3m4, _ => false.


Finding 3 — TensorLib/Dtype.lean:172 — typo in join comment

The new comment reads e4m3 has more range, e3m4 has more mantissa but niether dominate so no safe common type exists. Misspells "neither" and should be "dominates". Minor doc-quality cost only.


Finding 4 — TensorLib/Dtype.lean:1086 — missing double-rounding caveat for float64->e3m4

The | .float64, .float8_e3m4 arm rounds twice via Float32 (let f <- Float.ofLEByteArray data; return encodeFloat8E3M4 f.toFloat32) but lacks the double-rounding caveat comment the analogous | .float64, .float8_e4m3 arm carries on line 987 (rounds twice via fp32. Can disagree with ml_dtypes at the overflow edge (eg: 464.00000000000006)). Same failure mode applies at e3m4's own overflow edge (values just above 15.5 where fp64->fp32 rounds down to 15.5 exactly, then encodes finite instead of the single-round result). Cost: undocumented divergence from ml_dtypes for e3m4 while documented for e4m3.

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Fixes:

  • ToNpy now returns Err and rejects fp8_e3m4 with an error message instead of writing ambiguous V1 bytes.
  • Moved comment before fp8_e3m4 block.
  • Fixed typo in comment for join.
  • Added the double-rounding caveat comment matching the one on e4m3.

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Code review findings (single-pass, high-effort)

Note: single-pass review, Agent tool unavailable — no multi-agent fan-out or subagent verify pass ran.

Ranked most severe first.

1. TensorLib/Tensor.lean:682 — Breaking API signature change with no caller updates.
Tensor.toNpy changed from Tensor -> Npy.Ndarray to Tensor -> Err Npy.Ndarray. This is a silently breaking change for any external caller. The commented-out old body sits directly above (lines 673-680), so the intent was clearly deliberate, but the return type flip breaks source compatibility for let n := t.toNpy; n.save! path — every downstream call site must now handle the Err. In-repo Main.lean happens not to call Tensor.toNpy (it saves Ndarrays parsed from disk), so this went undetected.

2. TensorLib/Npy.lean:117 — Guard is at the wrong altitude; e3m4 -> npy corruption still reachable.
The check that blocks e3m4 was added in Tensor.toNpy (Tensor.lean:683), but the underlying serializer dtypeNameToNpyString still returns "V1" for .float8_e3m4, identical to .float8_e4m3. If a caller ever constructs an e3m4 Npy.Ndarray directly (e.g. Npy.Ndarray.mk { descr := { name := .float8_e3m4, ... }, ... } — the fields are public) and calls Ndarray.save!, Ndarray.toByteArray! writes a header tagged "V1", and Npy.parseFile reads it back as .float8_e4m3. Same bytes, different sign/exp/mantissa layout — silent numerical corruption. The right fix is at the serializer level: either dtypeNameToNpyString/toByteArray! returning Err on e3m4, or save! short-circuiting. As-is, this is a special-case bandaid on shared infrastructure.

3. TensorLib/Npy.lean:443 — Test pins the corruption behavior as a contract.
#guard Npy.Dtype.dtypeNameToNpyString .float8_e3m4 == "V1" locks in the buggy behavior from finding #2. If the serializer is later hardened (to reject e3m4 or use a distinct tag), this guard fails and could tempt someone into reverting the fix. Should instead assert that save! errors on e3m4, not that the wrong tag is produced.

4. TensorLib/Tensor.lean:673 — Commented-out old toNpy body left in place.
Lines 673-680 duplicate the new function body minus the guard. Git history is the record of what changed; leaving the dead copy makes future edits liable to resurrect the broken version or drift the two. Delete it.

5. TensorLib/Dtype.lean:137 — Constructor-qualifier inconsistency in isFloat.
The last alternative float8_e5m2 is unqualified while .float8_e4m3 and the newly added .float8_e3m4 are .-qualified. Cosmetic (Lean resolves both), and it's pre-existing, but the diff touches this exact line and it's the same pattern in isOneByte (line 100) and itemsize (line 145) — a one-line cleanup while the file is open.

No further correctness bugs survived a re-check against the diff. The e3m4 encode/decode routines mirror the e5m2 pattern that already has a PBT-verified round-trip, and the new e3m4 round-trip PBT is included. The joinOrdered additions preserve commutativity via the existing size-based swap (both dtypes are 1 byte, but the diff lists both directions e4m3,e3m4 and e3m4,e4m3 explicitly, so the commutativity PBT at line 1557 covers it).

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Fixes:

  • toNpy returns Err to make the e3m4 failure visible in the type system. No callers exist so nothing breaks internally. Added toNpy! as a wrapper for callers who don't need to handle the error.
  • Added an e3m4 guard in Ndarray.save! as well, so even direct Ndarray construction and save is blocked.
  • Can't directly test Ndarray.save! errors using a guard since it returns IO Unit (not pure). The fromNpyString "<V1" != .ok e3m4 guard already documents the limitation, and the runtime check in save! prevents the corruption.
  • Removed old function.
  • Fixed typo in isFloat.

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Single-pass code review (no Agent-tool fan-out). Findings:

1. TensorLib/Tensor.lean:673toNpy signature change is a silent API break.
toNpy went from Tensor -> Npy.Ndarray (total) to Tensor -> Err Npy.Ndarray. Any external caller using t.toNpy now fails to typecheck because they'd get an Err where they used to get an Ndarray. The diff adds toNpy! for the panic-on-error variant but doesn't migrate anyone or flag the break. In-repo there are no callers, so tests happily pass and hide the break.

2. TensorLib/Npy.lean:432 — duplicated e3m4 rejection at two layers.
Ndarray.save! re-checks header.descr.name == .float8_e3m4 even though Tensor.toNpy at Tensor.lean:674 already rejects e3m4 before an Ndarray can be constructed from a Tensor. Two guard sites for the same invariant, with slightly different error strings. Better altitude would be one rejection at the encoding boundary (e.g. inside dtypeNameToNpyString or Ndarray.toByteArray!). Cost: any future fp8 variant mapped to V1 needs both sites updated in lockstep, and a user constructing an Ndarray directly gets a different message than one going through Tensor.

3. TensorLib/Dtype.lean:188joinOrdered for e3m4 is inconsistent with e5m2.
The e3m4 arm mirrors e4m3's restrictive rules (only bool/int8/uint8/float32/float64), but e5m2 promotes int16/uint16/int32/uint32/int64/uint64/float16 to float32 or float64. Both e5m2 and e3m4 have similarly narrow dynamic range vs. those ints, so the divergence looks unintentional. join .float8_e3m4 .int16 returns none while join .float8_e5m2 .int16 returns float32. Either add matching float32/float64 promotions on the e3m4 side, or add a comment (like the one at line 172) explaining why e3m4 deliberately follows e4m3 rather than e5m2.

4. TensorLib/Tensor.lean:685 — indentation inconsistency on def toNpy!.
It's written with two leading spaces while adjacent top-level defs def ofNpy and def toNpy sit at column 0 inside the same namespace. Parses fine, purely cosmetic, but worth cleaning up while this code is fresh.

No correctness bugs found in the fp8_e3m4 encode/decode math itself — subnormal, normal, overflow-to-inf, NaN, and rounding boundaries all match the parallel e5m2 code and pass the ml_dtypes-verified #guard checks.

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Fixes:

  • There are no callers so nothing breaks internally. I added toNpy! for anyone who doesn't need to handle the error. Any external caller would get a compile error pointing them to the new signature, so it won't be silent.
  • The previous review pointed out that Tensor.toNpy alone can be bypassed by constructing an Ndarray directly, so the check exists at both layers. Moving it into dtypeNameToNpyString would require a bigger change for the same result.
  • This matches numpy. result_type(e3m4, int16) returns none like fp8_e4m3. fp8_e5m2 works differently here (because it's a float type while fp8_e4m3 / e3m4 are void types). Added a comment noting that e3m4 follows e4m3's rules.
  • Fixed indentation.

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Code review findings (single-pass, no Agent tool)

Single-pass inline review of the float8-e3m4 branch. No multi-agent fan-out / verify pass ran. The diff is small and well-tested (hermetic decode/encode #guards verified against ml_dtypes, PBTs for round-trip and commutativity, integration tests). No crash-class or wrong-answer bugs found — items below are API-shape / altitude / cleanup plus one join promotion gap. Ranked most-severe first.

1. TensorLib/Npy.lean:117 — e3m4 shares "V1" with e4m3 with no encoder-level guard.
dtypeNameToNpyString maps both float8_e3m4 and float8_e4m3 to "V1". The rejection lives in Tensor.toNpy and Ndarray.save!; the string encoder itself will happily produce a V1 header for e3m4. Any future serialization path that bypasses those two guards silently writes a file that reads back as e4m3, with completely different bit semantics and no error.

2. TensorLib/Tensor.lean:673toNpy signature is a breaking public API change.
Return type changed from Npy.Ndarray to Err Npy.Ndarray. No in-repo callers, but any external Lean project doing let n := arr.toNpy; n.save! path fails to type-check.

3. TensorLib/Npy.lean:431Ndarray.save! now throws, violating the ! convention.
In this codebase !-suffixed functions are total / panicking (toNpy!, isZero!, toByteArray!, and the new toNpy! added in the same diff). Turning save! into a throwing IO function silently changes its contract; callers that treat it as never-throwing get an uncaught IO exception on e3m4.

4. TensorLib/Dtype.lean:174joinOrdered for e3m4 is more restrictive than the comment claims.
The comment says "e3m4 follows e4m3's promotion rules" but the arms only cover float32/float64/bool/int8/uint8, falling through to none. e5m2 promotes with float16, int16/uint16/int32/uint32/int64/uint64; e3m4 does not. join(e3m4, float16) and join(e3m4, bfloat16) silently return none where numpy returns float32 / bfloat16. User writes tensor_e3m4 + tensor_float16 and gets a type-mismatch error rather than a promotion. Only the int16 divergence is documented.

5. TensorLib/Tensor.lean:674 & TensorLib/Npy.lean:432 — duplicated rejection with drifting error strings.
Tensor.toNpy says "...V1 which is indistinguishable from float8_e4m3"; Ndarray.save! says "...V1 which is identical to float8_e4m3". Two sources of truth for the same rule, already textually drifted. arr.toNpy!.save! panics at toNpy! before reaching save!, so the second check only fires for manually-constructed Ndarray.

6. TensorLib/Float.lean:340maxSafeNatForFloat8e3m4 = 15 comment is confusing.
Value is correct (16 overflows to +inf, so 15 is the largest integer with lossless round-trip), but the comment mixes "max value is 15.5" with an integer bound. The formula used for the other fp8 types (2^(m+1)) would give 32 here, which is wrong for e3m4 — the exponent range, not the mantissa, is the binding constraint. A future reviewer applying the formula uniformly could "correct" it.

7. TensorLib/Float.lean:508621 — three fp8 encoders are copy-paste with parameter drift.
toFloat8E4M3Bits, toFloat8E5M2Bits, and toFloat8E3M4Bits differ only in (expBits, mantBits, bias, overflowIsInf) plus a couple of magic thresholds (>=25 / >=26 / >=25, shift constants 14/7/17). RTE, subnormal shift, and overflow-carry are re-implemented three times; the next fp8 variant will make it four, and getting one of the shift constants wrong is easy.

Comment thread TensorLib/Tensor.lean
@seanmcl

seanmcl commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Reviewed this with a bit-exactness focus. First, the good news: the numeric core is exhaustively correct. I built the library clean (434 jobs, 0 errors / 0 warnings, so match exhaustiveness, the losslessAntiSymmetric theorem, and all #guards hold), then checked the codec bit-for-bit against independent oracles — all 256 decode bytes and ~34k encode sweep points (normal, subnormal, flush-to-zero, overflow, both signs), plus a cross-check against ml_dtypes 0.5.4. Zero mismatches. The encode/decode/arithmetic/cast/join/lossless math is solid.

There's exactly one genuine numeric divergence, plus some serialization/API/test items worth a look.

Should fix before merge

1. NaN encodes to the wrong byte — TensorLib/Float.lean:574 (confirmed vs ml_dtypes 0.5.4)

toFloat8E3M4Bits on a NaN returns sign8 ||| 0x71 (mantissa 0b0001 — the LSB). ml_dtypes emits 0x78 (mantissa 0b1000 — the quiet/MSB bit): fp32 NaN, 0/0, and inf - inf all give 0x78 (120) there but 0x71 (113) here. This also diverges from the sibling encoders in this file (e5m2 → 0x7E, e4m3 → 0x7F, both setting the quiet bit), and the adjacent comment -- quiet NaN (exp=7, mant=1) is itself off — mant=1 is the LSB, not the quiet bit. This breaks bit-exactness for any NaN-containing e3m4 array.

Suggested fix: use sign8 ||| 0x78, fix the comment, and add an encode #guard — e.g.

#guard (Float32.ofBits 0x7FC00000).toFloat8E3M4Bits == (0x78 : UInt8)

There's no encode #guard for NaN today, which is why this slipped through: the round-trip PBT's ∨ f != f disjunct makes every NaN input pass vacuously.

Highest-impact limitation (partly inherent — a design call)

2. npy "V1" collision silently corrupts on readTensorLib/Npy.lean:117,135

dtypeNameToNpyString maps both e3m4 and e4m3 to "V1", and fromNpyString "<V1" unconditionally returns .float8_e4m3. This PR guards the write path (toNpy / save!), but an e3m4 .npy produced externally (ml_dtypes/JAX also serialize e3m4 as V1) reads back as e4m3 with wrong values and no error — byte 0x6F is 15.5 in e3m4 vs 120.0 in e4m3. This is largely inherent to the format ambiguity, but nothing warns on the read side, and the write guard can read as "e3m4 + npy is handled" when the more common real-world path (loading foreign files) is still silently lossy. Might be worth an explicit ambiguity error or a note in the read path. (Also: the guards sit on the toNpy/save! wrappers, not on the byte emitter toByteArray!/toNpyString, so a hand-built e3m4 Ndarray reaching toByteArray! still emits a mislabeled buffer.)

API / design

3. toNpy! reintroduces the panic the refactor removes — TensorLib/Tensor.lean:685

toNpy was widened to Tensor → Err Npy.Ndarray (a breaking public-API change, made solely to reject e3m4; no in-repo caller). The new toNpy! := get! $ toNpy arr then panics on exactly the e3m4 case the Err refactor exists to prevent, and it has no callers — so it's a dead footgun. Consider dropping toNpy!, or guarding only save! and keeping toNpy total.

Test coverage

  • The serialization guards — this PR's entire e3m4 safety mechanism — have no tests (Test.lean). A later edit inverting either guard (==!=, or dropping the early .error) would ship green.
  • Untested new paths: NaN encode, subnormal encode (Float.lean:424-442, the most intricate code here), cross-fp8/fp16/bf16/fp64 casts, and negative overflow (-16 → -inf).
  • The round-trip PBT (Float.lean:803) is plausible-sampled (~100 of only 256 possible bytes) and sorry-admitted. Since the domain is just 256 values, an exhaustive #guard (List.range 256).all … (or native_decide) would be strictly stronger and deterministic.
  • The (-0) test (Test.lean:620) asserts v2 == 0.0, which +0.0 also satisfies — it doesn't actually verify sign preservation (would need v2.toBits == 0x80000000). This matches the existing fp16/e4m3 sibling tests, so it's a pre-existing convention rather than something new here.

Documented divergences (intentional — just flagging)

  • float64 → e3m4 double-rounds through fp32 and disagrees with ml_dtypes at interior values, not only "the overflow edge" the comment claims — e.g. fp64 1.03125… gives 1.0 here vs 1.0625 under single rounding. Inherited from the e4m3/e5m2 paths; the comment just understates the scope. (Dtype.lean:1086)
  • join(e3m4, e4m3) = none in both directions, where numpy promotes to e3m4. Documented as intentional; noting it as a numpy-parity divergence. (Dtype.lean:173)

Minor cleanups

  • Dead constant float8e3m4MantissaBits := 4 (Float.lean:38) is never referenced (maxSafeNatForFloat8e3m4 is hardcoded to 15). Worth deleting, or commenting why the sibling 2^(bits+1) formula doesn't apply — the bare 15 is a tamper hazard, since someone "fixing" it to 2^5 = 32 would silently accept 16..32 (which overflow to inf).
  • castOverflow fp8 duplication (Dtype.lean:1047): three near-identical ~130-line fp8 blocks. A decodeFloat8/encodeFloat8 dtype-dispatcher — analogous to the existing decodeFloat16OrBFloat16 — would collapse them.
  • Wrong comment -- 8 = maxSafeNat for e3m4 (Test.lean:642) — it's 15 (looks copy-pasted from the e4m3 test).
  • Divergent guard messages for one invariant: "indistinguishable from float8_e4m3" (Tensor.lean:674) vs "identical to float8_e4m3" (Npy.lean:432).
  • Cosmetic: 2-space comment indent at Dtype.lean:1047 where siblings use 4; and decode/encodeFloat8E3M4 lack the leading explanatory comment their e4m3/e5m2 siblings have.

Overall — nice work; the hard part (the bit math) is correct. #1 is the one I'd fix before merge.

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Fixes:

  1. Fixed — changed 0x71 to 0x78 so NaN encodes with the quiet bit (MSB of mantissa) set, matching ml_dtypes and the e4m3/e5m2 encoders. Added a #guard for +NaN encode. Skipped the -NaN guard because Lean normalizes all NaN to +NaN via Float32.toBits.
  2. Added warning comment in fromNpyString documenting that <V1 is ambiguous and defaults to e4m3. This is inherent to the format since there's no metadata to distinguish them on read. ToByteArray! is already private and its only call site is save!, which rejects e3m4 before reaching it. tldr; there is no code path where a hand written ndarray can reach this.
  3. Dropped !toNpy.

Test Coverage:

  1. Added #guard tests in Tensor.lean that assert toNpy returns .error for e3m4 and .ok for e4m3. If someone removes or inverts the guard, the build will break.
  2. Added #guard tests for NaN encode (0x78), subnormal encode (0.015625, 0.03125, 0.0625), and negative overflow (-16.0 -> -inf). Cross-fp8 casts are tested in Test.lean via castOverflow (e.g. fp32 --> e3m4, e3m4-->int8).
  3. Replaced the plausible-based PBT with an exhaustive #guard over all 256 values.
  4. Added a sign-preservation test for -0 that checks v2.toBits == 0x80000000 instead of relying on == 0.0.

Doc divergence:

Updated the double-rounding comment to reflect that the divergence applies at interior values, not just the overflow edge. join(e3m4, e4m3) = none is intentional; numpy promotes to e3m4 which silently overflows e4m3 values above 15.5.

Code cleanup:

  1. Deleted float8e3m4MantissaBits since it's unused. The comment on maxSafeNatForFloat8e3m4 explains why 15 is hardcoded.
  2. Refactored the 3 fp8 cast blocks into a shared helper to reduce code duplication.
  3. Fixed comment
  4. Fixed error msg wording to be consistent
  5. Fixed indentation and added comments explaining encode decode for fp8_e3m4

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Findings

1. Dead guard in Ndarray.save! (TensorLib/Npy.lean:433)
The arr.header.descr.name == .float8_e3m4 guard is unreachable through normal flows: parseFile/fromNpyString never produce an e3m4 descr (V1 always maps to e4m3), and Tensor.toNpy now blocks the write path before save! is reached. The only way to hit it is by hand-constructing an Ndarray. If it's intended as defense-in-depth, add a comment; otherwise remove — two copies of the error string are already drifting slightly between save! and toNpy.

2. Typos in Npy.lean comment (TensorLib/Npy.lean:134)
metadatmetadata, tyoestypes.

3. Misleading rationale in toNpy comment (TensorLib/Tensor.lean:675)
Comment says "Numpy has this limitation too," but numpy itself has no fp8 type — the V1 collision is an ml_dtypes/JAX artifact (they reuse the numpy void V1 slot for both fp8 subtypes). Clarify to avoid sending future readers on a numpy-side wild-goose chase.

4. Dead conditionals in e3m4 subnormal encode branch (TensorLib/Float.lean:604)
In the subnormal branch of Float32.toFloat8E3M4Bits, the guards if totalShift > 0 and if totalShift > 1 are unreachable — that branch is entered only when realExp < -2, forcing totalShift = 17 - realExp >= 20. This is copy-pasted from the e5m2/e4m3 encoders where the branch triggers at a boundary that makes those guards meaningful. Not a bug; drop the conditionals (or note they exist only for pattern consistency) so the reachable value range of totalShift is easier to reason about.

What I did not find

No correctness bugs in the encode/decode math: exponent rebias, mantissa shift widths (19 for e3m4 vs. 20/21 for e4m3/e5m2), subnormal scale factor (2^-6), overflow-to-inf boundary (realExp > 3), and RTNE round bit / sticky masks all check out against the E3M4 (1s + 3e + 4m, bias 3) spec. The exhaustive 256-value round-trip PBT and the arithmetic tests independently corroborate this.

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Fixes:

  1. Added a comment explaining it's normally unreachable but protects against hand-constructed Ndarrays.
  2. Fixed typos
  3. Used ml_dtypes instead of numpy.
  4. Fixed conditional.

@SmoothThunk

Copy link
Copy Markdown
Collaborator Author

Re-review results

Fix commit fa99aab addresses all four prior findings:

  • F1 (dead guard): kept as defense-in-depth with an explanatory comment at Npy.lean:433-434 explaining unreachability and intent. Good.
  • F2 (typos): metadatmetadata, tyoestypes at Npy.lean:134. Correct.
  • F3 (misleading rationale): "Numpy has this limitation too" → "Ml_dtypes has this limitation too" at Tensor.lean:675. Correct (minor nit — capitalized as "Ml_dtypes"; canonical spelling is ml_dtypes, but this is style-level).
  • F4 (dead conditionals): removed the two unreachable branches at Float.lean:611-613 and added a comment -- totalShift >= 20 in this branch (realExp < -2), so always > 1. The reachability claim is correct: entry to this branch requires realExp < -2, so totalShift = 17 - realExp >= 20, and 20 - 1 = 19 < 32 keeps the shift within UInt32 bounds. The stickyMask shift 1 <<< (totalShift - 1) also stays safe for the totalShift <= 24 subrange (the outer guard totalShift >= 25 returns early). Correct.

Fresh pass over the whole main...HEAD diff produced no new correctness or cleanup issues beyond what was already reported and fixed.

@SmoothThunk
SmoothThunk merged commit af390ae into leanprover:main Jul 28, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants